#include <iostream>
#include <fstream>
#include <cstring>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <chrono>

#define PORT 8081
#define BUFFER_SIZE 65507  // UDP最大数据报大小

int main() {
    int sockfd;
    struct sockaddr_in servaddr, cliaddr;
    char buffer[BUFFER_SIZE] = {0};
    socklen_t len;
    const char* filename = "received_udp_file.bin";
    
    // 创建套接字
    if ((sockfd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("socket creation failed");
        exit(EXIT_FAILURE);
    }
    
    memset(&servaddr, 0, sizeof(servaddr));
    memset(&cliaddr, 0, sizeof(cliaddr));
    
    // 设置服务器地址
    servaddr.sin_family = AF_INET;
    servaddr.sin_addr.s_addr = INADDR_ANY;
    servaddr.sin_port = htons(PORT);
    
    // 绑定套接字
    if (bind(sockfd, (const struct sockaddr *)&servaddr, sizeof(servaddr)) < 0) {
        perror("bind failed");
        exit(EXIT_FAILURE);
    }
    
    std::cout << "UDP Server listening on port " << PORT << std::endl;
    
    std::ofstream file(filename, std::ios::binary);
    if (!file) {
        perror("file open failed");
        exit(EXIT_FAILURE);
    }
    
    auto start_time = std::chrono::high_resolution_clock::now();
    size_t total_bytes = 0;
    len = sizeof(cliaddr);
    
    // 接收文件
    while (true) {
        int bytes_received = recvfrom(sockfd, (char *)buffer, BUFFER_SIZE, 
                                      MSG_WAITALL, (struct sockaddr *)&cliaddr, &len);
        if (bytes_received <= 0) {
            break;
        }
        
        file.write(buffer, bytes_received);
        total_bytes += bytes_received;
        
        // 发送确认（UDP需要应用层确认机制）
        sendto(sockfd, "ACK", 3, MSG_CONFIRM, 
               (const struct sockaddr *)&cliaddr, len);
    }
    
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::seconds>(end_time - start_time).count();
    
    file.close();
    close(sockfd);
    
    std::cout << "UDP File received successfully: " << total_bytes << " bytes in " 
              << duration << " seconds" << std::endl;
    if (duration > 0) {
        std::cout << "Effective transfer rate: " 
                  << (total_bytes / (1024.0 * 1024.0)) / duration << " MB/s" << std::endl;
    }
    
    return 0;
}    